#!/usr/bin/python
# -*- coding:utf-8 -*-
# @FileName : DL5_test6_1.py
# Author    : myh
import torch
from torch import nn
import time

def try_gpu(i=0):  #@save
    """如果存在，则返回gpu(i)，否则返回cpu()"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def try_all_gpus():  #@save
    """返回所有可用的GPU，如果没有GPU，则返回[cpu(),]"""
    devices = [torch.device(f'cuda:{i}')
             for i in range(torch.cuda.device_count())]
    return devices if devices else [torch.device('cpu')]

# print(try_gpu(), try_gpu(10), try_all_gpus())
# x = torch.tensor([1, 2, 3])
# print(x.device)
#
#
# X = torch.ones(2, 3, device=try_gpu())
# print(X)
# print(torch.device('cpu'), torch.device('cuda'), torch.device('cuda:1'))
# print(torch.cuda.device_count())

class Timer:
    def __init__(self):
        self.times = []
        self.start()


    def start(self):
        self.start_time = time.time()

    def stop(self):
        self.times.append(time.time() - self.start_time)
        return self.times[-1]

    def avg(self):
        return sum(self.times)/len(self.times)

    def sum(self):
        return sum(self.times)

    def cumsum(self):
        return np.array(self.times).cumsum().tolist()


#
# time1 = Timer()
# X = torch.normal(mean=0.5*torch.ones(20000,30000), std=1).cuda()
# Y = torch.normal(mean=0.5*torch.ones(30000,20000), std=1).cuda()
# Z = torch.mm(X,Y)
# # print(Z)
# print(f"{time1.stop()}")
#
#
# time2 = Timer()
# X = torch.normal(mean=0.5*torch.ones(20000,30000), std=1)
# Y = torch.normal(mean=0.5*torch.ones(30000,20000), std=1)
# Z = torch.mm(X,Y)
# # print(Z)
# print(f"{time2.stop()}")



A = torch.rand(100, 100, device=try_gpu())

start = time.time()
for i in range(1000):
    A = torch.mm(A, A)
    B = torch.norm(A)  # 逐个记录
end = time.time()
print(f'逐个记录耗费时间：{round((end - start) * 1000)} ms')

A = torch.rand(100, 100, device=try_gpu())
start = time.time()
for i in range(1000):
    A = torch.mm(A, A)
B = torch.norm(A)  # 最终记录
end = time.time()
print(f'最终记录耗费时间：{round((end - start) * 1000)} ms')
